import sys
import akshare as ak
from PyQt5 import QtCore, QtWidgets, QtGui
from PyQt5.QtCore import Qt
from datetime import datetime

class StockWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.load_data()

    def initUI(self):
        self.setGeometry(300, 300, 800, 600)
        self.setWindowTitle("上证指数K线图 - 日线")
        # 添加立体边框样式
        self.setStyleSheet("""
            QMainWindow {
                background-color: #3A3A3A;
                border: 2px solid #C0C0C0;
                border-radius: 5px;
            }
            QStatusBar {
                color: #B0B0B0;
                background: #404040;
            }
        """)
        
        # 先创建 scene 再创建 view
        self.scene = QtWidgets.QGraphicsScene()  # 添加这行
        
        # 创建带阴影的视图
        self.view = QtWidgets.QGraphicsView(self.scene)  # 这行保持原位
        self.view.setStyleSheet("""
            QGraphicsView {
                border: 1px solid #808080;
                border-radius: 3px;
                background: #262626;
            }
        """)
        # 隐藏滚动条
        self.view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        
        # 添加操作提示标签
        self.info_label = QtWidgets.QLabel("↑ 放大  ↓ 缩小  ← 左移  → 右移")
        self.info_label.setAlignment(Qt.AlignCenter)
        self.info_label.setStyleSheet("""
            QLabel {
                background: #404040;
                color: #C0C0C0;
                padding: 8px;
                border: 1px solid #606060;
                border-radius: 8px;
                font: bold 14px;
            }
        """)
        self.info_label.setGeometry(300, 500, 200, 40)
        self.scene.addWidget(self.info_label)

        # 添加坐标变换（确保Y轴方向正确）
        transform = QtGui.QTransform()
        transform.scale(1, -1)
        self.view.setTransform(transform)  # 恢复Y轴翻转
        
        # 修正坐标系转换（关键修改）
        transform = QtGui.QTransform()
        transform.scale(1, -1)
        transform.translate(0, -800)  # 匹配绘图区域高度
        self.view.setTransform(transform)

    def load_data(self):  # 添加这个方法
        # 测试数据
        self.chart_data = [{
            'open': 3000 + i,
            'high': 3020 + i,
            'low': 2980 + i,
            'close': 3010 + i
        } for i in range(100)]
        self.draw_chart()

    def draw_chart(self):
        self.scene.clear()
        item = CandlestickItem(self.chart_data)
        self.scene.addItem(item)
        # 确保视图正确显示
        self.view.fitInView(item.boundingRect(), Qt.KeepAspectRatio)

class CandlestickItem(QtWidgets.QGraphicsItem):
    def __init__(self, data, parent=None):  # 添加 parent 参数
        super().__init__(parent)  # 传递父项参数
        self.data = data
        self.rect_width = 6
        self.min_price = min(row['low'] for row in self.data)
        self.max_price = max(row['high'] for row in self.data)
        self.price_range = self.max_price - self.min_price

    def boundingRect(self):
        # 增加高度保证可见
        return QtCore.QRectF(0, 0, (len(self.data)+30)*self.rect_width, 800)

    def paint(self, painter, option, widget):
        view_height = 800  # 必须与 boundingRect 高度一致
        
        for i, row in enumerate(self.data):
            # 修改影线绘制方式（添加颜色）
            painter.setPen(QtGui.QPen(QtGui.QColor(150, 150, 150), 1))
            painter.drawLine(x + self.rect_width/2, high_y, 
                            x + self.rect_width/2, low_y)
            
            # 修改实体绘制逻辑
            body_height = abs(close_y - open_y)
            if body_height < 1:  # 处理极窄实体
                body_height = 1
            painter.drawRect(x, min(open_y, close_y), 
                            self.rect_width, body_height)

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    win = StockWindow()
    win.show()
    sys.exit(app.exec_())
